// Time:  O(m * n)
// Space: O(1)

#include <vector>
#include <iostream>
#include <cassert>

using namespace std;

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& matrix) {
        // （0，0） -> （0，1） -> （1，0） -> （2，0） -> （1，1） -> （0，2） -> （1，2） -> （2，1） -> （3，0）
        if(matrix.empty()) {
            return {};
        }

        const int M = matrix.size();
        const int N = matrix[0].size();

        vector<int> result;

        for(int s = 0; s < M + N - 1; s++) {
            // for all i + j = s

            for(int x = 0; x <= s; x++) {
                int i = x;
                int j = s - x;
                if(s % 2 == 0) {
                    swap(i, j);
                }

                if(i >= M || j >= N) {
                    continue;
                }
                result.push_back(matrix[i][j]);
            }
        }
        
        return result;
    }
};

int main() {
    Solution solution;

    vector<vector<int>> matrix = {
        { 1, 2, 3 },
        { 4, 5, 6 },
        { 7, 8, 9 }
    };

    vector<int> expected = { 1, 2, 4, 7, 5, 3, 6, 8, 9 };
    vector<int> result = solution.findDiagonalOrder(matrix);

    // compare the result and the expected
    assert(result == expected);

    vector<vector<int>> matrix1 = {};

    vector<int> expected1 = {};
    vector<int> result1 = solution.findDiagonalOrder(matrix1);

    assert(result1 == expected1);

    vector<vector<int>> matrix2 = {
        { 1, 2, 3, 4 }
    };

    vector<int> expected2 = { 1, 2, 3, 4 };
    vector<int> result2 = solution.findDiagonalOrder(matrix2);

    assert(result2 == expected2);

    vector<vector<int>> matrix3 = {
        { 1 },
        { 2 },
        { 3 },
        { 4 }
    };

    vector<int> expected3 = { 1, 2, 3, 4 };
    vector<int> result3 = solution.findDiagonalOrder(matrix3);

    assert(result3 == expected3);
}
